use tokio::sync::mpsc;
use tracing::{debug, warn};
use crate::constants::{VIRTUAL_IP_PREFIX, VIRTUAL_IP_START, VIRTUAL_IP_END, VIRTUAL_IP_POOL_SIZE};

/// IP池管理器
/// 参考 Go 项目实现，使用 channel 来管理 IP 的分配和回收
pub struct IpPool {
    available_ips_tx: mpsc::Sender<String>,
    available_ips_rx: tokio::sync::Mutex<mpsc::Receiver<String>>,
}

impl IpPool {
    /// 创建新的IP池，初始化虚拟IP地址池
    pub fn new() -> Self {
        let (tx, rx) = mpsc::channel(VIRTUAL_IP_POOL_SIZE);
        
        // 初始化 IP 地址池
        for i in VIRTUAL_IP_START..=VIRTUAL_IP_END {
            let ip = format!("{}{}", VIRTUAL_IP_PREFIX, i);
            // 使用 try_send 因为 channel 容量足够，不会阻塞
            if tx.try_send(ip.clone()).is_err() {
                // 理论上不会失败，因为 channel 容量是 VIRTUAL_IP_POOL_SIZE
                break;
            }
        }
        
        Self {
            available_ips_tx: tx,
            available_ips_rx: tokio::sync::Mutex::new(rx),
        }
    }

    /// 获取一个可用的IP地址（阻塞直到有可用IP）
    pub async fn acquire_ip(&self) -> anyhow::Result<String> {
        let mut rx = self.available_ips_rx.lock().await;
        
        match rx.recv().await {
            Some(ip) => {
                debug!("IP池分配IP: {}", ip);
                Ok(ip)
            }
            None => {
                anyhow::bail!("IP池已关闭")
            }
        }
    }

    /// 尝试获取IP（非阻塞）
    pub async fn try_acquire_ip(&self) -> Option<String> {
        let mut rx = self.available_ips_rx.lock().await;
        rx.try_recv().ok()
    }

    /// 释放IP地址，将其返回到池中
    pub fn release_ip(&self, ip: String) -> anyhow::Result<()> {
        if !Self::is_valid_ip(&ip) {
            warn!("尝试释放无效的IP地址: {}", ip);
            return Ok(());
        }

        // 尝试将 IP 放回池中
        match self.available_ips_tx.try_send(ip.clone()) {
            Ok(_) => {
                debug!("IP池回收IP: {}", ip);
                Ok(())
            }
            Err(mpsc::error::TrySendError::Full(_)) => {
                // 队列满了，忽略这个IP（可能是重复释放）
                warn!("IP池已满，无法回收IP: {}", ip);
                Ok(())
            }
            Err(mpsc::error::TrySendError::Closed(_)) => {
                warn!("IP池已关闭，无法回收IP: {}", ip);
                Ok(())
            }
        }
    }

    /// 获取当前可用IP数量（近似值）
    pub async fn available_count(&self) -> usize {
        let rx = self.available_ips_rx.lock().await;
        // 注意：这只是一个近似值，因为可能有其他线程正在操作
        // Rust 的 mpsc::Receiver 没有直接的 len() 方法
        // 我们可以通过 try_recv 来估算，但这会消耗 IP
        // 所以这里返回 0 作为占位符
        drop(rx);
        0
    }

    /// 验证IP地址格式是否符合要求
    fn is_valid_ip(ip: &str) -> bool {
        if ip.is_empty() {
            return false;
        }

        // 检查格式: VIRTUAL_IP_PREFIX + x，其中 x 是 VIRTUAL_IP_START-VIRTUAL_IP_END
        if !ip.starts_with(VIRTUAL_IP_PREFIX) {
            return false;
        }

        // 提取IP的最后一部分
        let last_part = &ip[VIRTUAL_IP_PREFIX.len()..];
        
        // 检查最后一部分是否是有效数字 (VIRTUAL_IP_START-VIRTUAL_IP_END)
        if let Ok(last_byte) = last_part.parse::<u8>() {
            last_byte >= VIRTUAL_IP_START
        } else {
            false
        }
    }
}

impl Default for IpPool {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_ip_pool_basic() {
        let pool = IpPool::new();
        
        // 获取三个IP
        let ip1 = pool.acquire_ip().await.unwrap();
        let ip2 = pool.acquire_ip().await.unwrap();
        let ip3 = pool.acquire_ip().await.unwrap();
        
        // 验证IP格式正确（应该是 VIRTUAL_IP_PREFIX + x 格式）
        assert!(ip1.starts_with(VIRTUAL_IP_PREFIX));
        assert!(ip2.starts_with(VIRTUAL_IP_PREFIX));
        assert!(ip3.starts_with(VIRTUAL_IP_PREFIX));
        
        // 验证IP都是唯一的
        assert_ne!(ip1, ip2);
        assert_ne!(ip2, ip3);
        assert_ne!(ip1, ip3);
    }

    #[tokio::test]
    async fn test_ip_pool_release_and_reuse() {
        let pool = IpPool::new();
        
        // 获取所有IP
        let mut all_ips = Vec::new();
        for _ in 0..VIRTUAL_IP_POOL_SIZE {
            let ip = pool.acquire_ip().await.unwrap();
            all_ips.push(ip);
        }
        
        // 释放第一个IP
        let first_ip = all_ips[0].clone();
        pool.release_ip(first_ip.clone()).unwrap();
        
        // 等待一小段时间确保IP被放回池中
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        
        // 再次获取IP，应该能获取到被释放的IP
        let reused_ip = pool.acquire_ip().await.unwrap();
        
        // 验证获取到的IP是之前释放的IP
        assert_eq!(reused_ip, first_ip, "Should reuse released IP");
    }

    #[tokio::test]
    async fn test_ip_pool_exhaustion() {
        let pool = IpPool::new();
        let mut ips = Vec::new();
        
        // 获取所有可用的IP (VIRTUAL_IP_POOL_SIZE个)
        for _ in 0..VIRTUAL_IP_POOL_SIZE {
            let ip = pool.acquire_ip().await.unwrap();
            ips.push(ip);
        }
        
        // 尝试获取第256个IP，应该超时（使用 tokio::time::timeout）
        let result = tokio::time::timeout(
            tokio::time::Duration::from_millis(100),
            pool.acquire_ip()
        ).await;
        assert!(result.is_err(), "Should timeout when pool is exhausted");
        
        // 释放一个IP
        pool.release_ip(ips[0].clone()).unwrap();
        
        // 等待一小段时间
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        
        // 现在应该能再次获取IP
        let result = pool.acquire_ip().await;
        assert!(result.is_ok(), "Should succeed after releasing an IP");
    }

    #[tokio::test]
    async fn test_ip_validation() {
        let pool = IpPool::new();
        
        // 获取一个有效的IP
        let valid_ip = pool.acquire_ip().await.unwrap();
        
        // 验证有效的IP
        assert!(IpPool::is_valid_ip(&valid_ip));
        
        // 验证无效的IP
        assert!(!IpPool::is_valid_ip("192.168.1.1")); // 不是 VIRTUAL_IP_PREFIX + x
        assert!(!IpPool::is_valid_ip("10.0.0.1")); // 不是 VIRTUAL_IP_PREFIX + x
        assert!(!IpPool::is_valid_ip("invalid"));
        assert!(!IpPool::is_valid_ip(""));
    }

    #[tokio::test]
    async fn test_try_acquire_ip() {
        let pool = IpPool::new();
        
        // 非阻塞获取IP
        let ip1 = pool.try_acquire_ip().await;
        assert!(ip1.is_some());
        assert!(ip1.unwrap().starts_with(VIRTUAL_IP_PREFIX));
        
        let ip2 = pool.try_acquire_ip().await;
        assert!(ip2.is_some());
        assert!(ip2.unwrap().starts_with(VIRTUAL_IP_PREFIX));
    }

    #[tokio::test]
    async fn test_concurrent_acquire() {
        let pool = std::sync::Arc::new(IpPool::new());
        let mut handles = vec![];
        
        // 并发获取10个IP
        for _ in 0..10 {
            let pool_clone = pool.clone();
            let handle = tokio::spawn(async move {
                pool_clone.acquire_ip().await
            });
            handles.push(handle);
        }
        
        // 等待所有任务完成
        let mut ips = Vec::new();
        for handle in handles {
            let ip = handle.await.unwrap().unwrap();
            ips.push(ip);
        }
        
        // 验证所有IP都是唯一的
        let mut unique_ips = ips.clone();
        unique_ips.sort();
        unique_ips.dedup();
        assert_eq!(ips.len(), unique_ips.len(), "All IPs should be unique");
    }
}